{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LongRAG Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to implement LongRAG using LlamaIndex workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/user-attachments/files/16474262/data.zip -O data.zip\n",
    "!unzip -o data.zip\n",
    "!rm data.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since workflows are async first, this all runs fine in a notebook. If you were running in your own code, you would want to use `asyncio.run()` to start an async event loop if one isn't already running.\n",
    "\n",
    "```python\n",
    "async def main():\n",
    "    <async code>\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import asyncio\n",
    "    asyncio.run(main())\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper Functions\n",
    "\n",
    "These helper functions will help us split documents into smaller pieces and group  nodes based on their relationships."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Dict, Optional, Set, FrozenSet\n",
    "\n",
    "from llama_index.core.schema import BaseNode, TextNode\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "# constants\n",
    "DEFAULT_CHUNK_SIZE = 4096  # optionally splits documents into CHUNK_SIZE, then regroups them to demonstrate grouping algorithm\n",
    "DEFAULT_MAX_GROUP_SIZE = 20  # maximum number of documents in a group\n",
    "DEFAULT_SMALL_CHUNK_SIZE = 512  # small chunk size for generating embeddings\n",
    "DEFAULT_TOP_K = 8  # top k for retrieving\n",
    "\n",
    "\n",
    "def split_doc(chunk_size: int, documents: List[BaseNode]) -> List[TextNode]:\n",
    "    \"\"\"Splits documents into smaller pieces.\n",
    "\n",
    "    Args:\n",
    "        chunk_size (int): Chunk size\n",
    "        documents (List[BaseNode]): Documents\n",
    "\n",
    "    Returns:\n",
    "        List[TextNode]: Smaller chunks\n",
    "    \"\"\"\n",
    "    # split docs into tokens\n",
    "    text_parser = SentenceSplitter(chunk_size=chunk_size)\n",
    "    return text_parser.get_nodes_from_documents(documents)\n",
    "\n",
    "\n",
    "def group_docs(\n",
    "    nodes: List[str],\n",
    "    adj: Dict[str, List[str]],\n",
    "    max_group_size: Optional[int] = DEFAULT_MAX_GROUP_SIZE,\n",
    ") -> Set[FrozenSet[str]]:\n",
    "    \"\"\"Groups documents.\n",
    "\n",
    "    Args:\n",
    "        nodes (List[str]): documents IDs\n",
    "        adj (Dict[str, List[str]]): related documents for each document; id -> list of doc strings\n",
    "        max_group_size (Optional[int], optional): max group size, None if no max group size. Defaults to DEFAULT_MAX_GROUP_SIZE.\n",
    "    \"\"\"\n",
    "    docs = sorted(nodes, key=lambda node: len(adj[node]))\n",
    "    groups = set()  # set of set of IDs\n",
    "    for d in docs:\n",
    "        related_groups = set()\n",
    "        for r in adj[d]:\n",
    "            for g in groups:\n",
    "                if r in g:\n",
    "                    related_groups = related_groups.union(frozenset([g]))\n",
    "\n",
    "        gnew = {d}\n",
    "        related_groupsl = sorted(related_groups, key=lambda el: len(el))\n",
    "        for g in related_groupsl:\n",
    "            if max_group_size is None or len(gnew) + len(g) <= max_group_size:\n",
    "                gnew = gnew.union(g)\n",
    "                if g in groups:\n",
    "                    groups.remove(g)\n",
    "\n",
    "        groups.add(frozenset(gnew))\n",
    "\n",
    "    return groups\n",
    "\n",
    "\n",
    "def get_grouped_docs(\n",
    "    nodes: List[TextNode],\n",
    "    max_group_size: Optional[int] = DEFAULT_MAX_GROUP_SIZE,\n",
    ") -> List[TextNode]:\n",
    "    \"\"\"Gets list of documents that are grouped.\n",
    "\n",
    "    Args:\n",
    "        nodes (t.List[TextNode]): Input list\n",
    "        max_group_size (Optional[int], optional): max group size, None if no max group size. Defaults to DEFAULT_MAX_GROUP_SIZE.\n",
    "\n",
    "    Returns:\n",
    "        t.List[TextNode]: Output list\n",
    "    \"\"\"\n",
    "    # node IDs\n",
    "    nodes_str = [node.id_ for node in nodes]\n",
    "    # maps node ID -> related node IDs based on that node's relationships\n",
    "    adj: Dict[str, List[str]] = {\n",
    "        node.id_: [val.node_id for val in node.relationships.values()]\n",
    "        for node in nodes\n",
    "    }\n",
    "    # node ID -> node\n",
    "    nodes_dict = {node.id_: node for node in nodes}\n",
    "\n",
    "    res = group_docs(nodes_str, adj, max_group_size)\n",
    "\n",
    "    ret_nodes = []\n",
    "    for g in res:\n",
    "        cur_node = TextNode()\n",
    "\n",
    "        for node_id in g:\n",
    "            cur_node.text += nodes_dict[node_id].text + \"\\n\\n\"\n",
    "            cur_node.metadata.update(nodes_dict[node_id].metadata)\n",
    "\n",
    "        ret_nodes.append(cur_node)\n",
    "\n",
    "    return ret_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making the Retriever\n",
    "\n",
    "LongRAG needs a custom retriever, which is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.retrievers import BaseRetriever\n",
    "from llama_index.core.vector_stores.simple import BasePydanticVectorStore\n",
    "from llama_index.core.schema import QueryBundle, NodeWithScore\n",
    "from llama_index.core.vector_stores.types import VectorStoreQuery\n",
    "from llama_index.core.settings import Settings\n",
    "\n",
    "\n",
    "class LongRAGRetriever(BaseRetriever):\n",
    "    \"\"\"Long RAG Retriever.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        grouped_nodes: List[TextNode],\n",
    "        small_toks: List[TextNode],\n",
    "        vector_store: BasePydanticVectorStore,\n",
    "        similarity_top_k: int = DEFAULT_TOP_K,\n",
    "    ) -> None:\n",
    "        \"\"\"Constructor.\n",
    "\n",
    "        Args:\n",
    "            grouped_nodes (List[TextNode]): Long retrieval units, nodes with docs grouped together based on relationships\n",
    "            small_toks (List[TextNode]): Smaller tokens\n",
    "            embed_model (BaseEmbedding, optional): Embed model. Defaults to None.\n",
    "            similarity_top_k (int, optional): Similarity top k. Defaults to 8.\n",
    "        \"\"\"\n",
    "        self._grouped_nodes = grouped_nodes\n",
    "        self._grouped_nodes_dict = {node.id_: node for node in grouped_nodes}\n",
    "        self._small_toks = small_toks\n",
    "        self._small_toks_dict = {node.id_: node for node in self._small_toks}\n",
    "\n",
    "        self._similarity_top_k = similarity_top_k\n",
    "        self._vec_store = vector_store\n",
    "        self._embed_model = Settings.embed_model\n",
    "\n",
    "    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:\n",
    "        \"\"\"Retrieves.\n",
    "\n",
    "        Args:\n",
    "            query_bundle (QueryBundle): query bundle\n",
    "\n",
    "        Returns:\n",
    "            List[NodeWithScore]: nodes with scores\n",
    "        \"\"\"\n",
    "        # make query\n",
    "        query_embedding = self._embed_model.get_query_embedding(\n",
    "            query_bundle.query_str\n",
    "        )\n",
    "        vector_store_query = VectorStoreQuery(\n",
    "            query_embedding=query_embedding, similarity_top_k=500\n",
    "        )\n",
    "\n",
    "        # query for answer\n",
    "        query_res = self._vec_store.query(vector_store_query)\n",
    "\n",
    "        # determine top parents of most similar children (these are long retrieval units)\n",
    "        top_parents_set: Set[str] = set()\n",
    "        top_parents: List[NodeWithScore] = []\n",
    "        for id_, similarity in zip(query_res.ids, query_res.similarities):\n",
    "            cur_node = self._small_toks_dict[id_]\n",
    "            parent_id = cur_node.ref_doc_id\n",
    "            if parent_id not in top_parents_set:\n",
    "                top_parents_set.add(parent_id)\n",
    "\n",
    "                parent_node = self._grouped_nodes_dict[parent_id]\n",
    "                node_with_score = NodeWithScore(\n",
    "                    node=parent_node, score=similarity\n",
    "                )\n",
    "                top_parents.append(node_with_score)\n",
    "\n",
    "                if len(top_parents_set) >= self._similarity_top_k:\n",
    "                    break\n",
    "\n",
    "        assert len(top_parents) == min(\n",
    "            self._similarity_top_k, len(self._grouped_nodes)\n",
    "        )\n",
    "\n",
    "        return top_parents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Designing the Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LongRAG consists of the following steps:\n",
    "\n",
    "1. Ingesting the data — grouping documents and putting them in long retrieval units, splitting the long retrieval units into smaller tokens to generate embeddings, and indexing the small nodes.\n",
    "2. Constructing the retriever and query engine.\n",
    "3. Querying over the data given a string."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define an event that passes the long and small retrieval units into the retriever and query engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterable\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.llms import LLM\n",
    "from llama_index.core.workflow import Event\n",
    "\n",
    "\n",
    "class LoadNodeEvent(Event):\n",
    "    \"\"\"Event for loading nodes.\"\"\"\n",
    "\n",
    "    small_nodes: Iterable[TextNode]\n",
    "    grouped_nodes: list[TextNode]\n",
    "    index: VectorStoreIndex\n",
    "    similarity_top_k: int\n",
    "    llm: LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After defining our events, we can write our workflow and steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Workflow,\n",
    "    step,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Context,\n",
    ")\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "\n",
    "\n",
    "class LongRAGWorkflow(Workflow):\n",
    "    \"\"\"Long RAG Workflow.\"\"\"\n",
    "\n",
    "    @step\n",
    "    async def ingest(self, ev: StartEvent) -> LoadNodeEvent | None:\n",
    "        \"\"\"Ingestion step.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): Context\n",
    "            ev (StartEvent): start event\n",
    "\n",
    "        Returns:\n",
    "            StopEvent | None: stop event with result\n",
    "        \"\"\"\n",
    "        data_dir: str = ev.get(\"data_dir\")\n",
    "        llm: LLM = ev.get(\"llm\")\n",
    "        chunk_size: int | None = ev.get(\"chunk_size\")\n",
    "        similarity_top_k: int = ev.get(\"similarity_top_k\")\n",
    "        small_chunk_size: int = ev.get(\"small_chunk_size\")\n",
    "        index: VectorStoreIndex | None = ev.get(\"index\")\n",
    "        index_kwargs: dict[str, t.Any] | None = ev.get(\"index_kwargs\")\n",
    "\n",
    "        if any(\n",
    "            i is None\n",
    "            for i in [data_dir, llm, similarity_top_k, small_chunk_size]\n",
    "        ):\n",
    "            return None\n",
    "\n",
    "        if not index:\n",
    "            docs = SimpleDirectoryReader(data_dir).load_data()\n",
    "            if chunk_size is not None:\n",
    "                nodes = split_doc(\n",
    "                    chunk_size, docs\n",
    "                )  # split documents into chunks of chunk_size\n",
    "                grouped_nodes = get_grouped_docs(\n",
    "                    nodes\n",
    "                )  # get list of nodes after grouping (groups are combined into one node), these are long retrieval units\n",
    "            else:\n",
    "                grouped_nodes = docs\n",
    "\n",
    "            # split large retrieval units into smaller nodes\n",
    "            small_nodes = split_doc(small_chunk_size, grouped_nodes)\n",
    "\n",
    "            index_kwargs = index_kwargs or {}\n",
    "            index = VectorStoreIndex(small_nodes, **index_kwargs)\n",
    "        else:\n",
    "            # get smaller nodes from index and form large retrieval units from these nodes\n",
    "            small_nodes = index.docstore.docs.values()\n",
    "            grouped_nodes = get_grouped_docs(small_nodes, None)\n",
    "\n",
    "        return LoadNodeEvent(\n",
    "            small_nodes=small_nodes,\n",
    "            grouped_nodes=grouped_nodes,\n",
    "            index=index,\n",
    "            similarity_top_k=similarity_top_k,\n",
    "            llm=llm,\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def make_query_engine(\n",
    "        self, ctx: Context, ev: LoadNodeEvent\n",
    "    ) -> StopEvent:\n",
    "        \"\"\"Query engine construction step.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): context\n",
    "            ev (LoadNodeEvent): event\n",
    "\n",
    "        Returns:\n",
    "            StopEvent: stop event\n",
    "        \"\"\"\n",
    "        # make retriever and query engine\n",
    "        retriever = LongRAGRetriever(\n",
    "            grouped_nodes=ev.grouped_nodes,\n",
    "            small_toks=ev.small_nodes,\n",
    "            similarity_top_k=ev.similarity_top_k,\n",
    "            vector_store=ev.index.vector_store,\n",
    "        )\n",
    "        query_eng = RetrieverQueryEngine.from_args(retriever, ev.llm)\n",
    "\n",
    "        return StopEvent(\n",
    "            result={\n",
    "                \"retriever\": retriever,\n",
    "                \"query_engine\": query_eng,\n",
    "                \"index\": ev.index,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def query(self, ctx: Context, ev: StartEvent) -> StopEvent | None:\n",
    "        \"\"\"Query step.\n",
    "\n",
    "        Args:\n",
    "            ctx (Context): context\n",
    "            ev (StartEvent): start event\n",
    "\n",
    "        Returns:\n",
    "            StopEvent | None: stop event with result\n",
    "        \"\"\"\n",
    "        query_str: str | None = ev.get(\"query_str\")\n",
    "        query_eng = ev.get(\"query_eng\")\n",
    "\n",
    "        if query_str is None:\n",
    "            return None\n",
    "\n",
    "        result = query_eng.query(query_str)\n",
    "        return StopEvent(result=result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Walkthrough:\n",
    "- There are 2 entry points: one for ingesting and indexing and another for querying.\n",
    "- When ingesting, it first reads the documents, splits them into smaller nodes, and indexes them. After that, it sends a `LoadNodeEvent` which triggers the execution of `make_query_engine`, constructing a retriever and query engine from the nodes. It returns a result of the retriever, the query engine, and the index.\n",
    "- When querying, it takes in the query from the `StartEvent`, feeds it into the query engine in the context, and returns the result of the query.\n",
    "- The context is used to store the query engine."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "wf = LongRAGWorkflow(timeout=60)\n",
    "llm = OpenAI(\"gpt-4o\")\n",
    "data_dir = \"data\"\n",
    "\n",
    "# initialize the workflow\n",
    "result = await wf.run(\n",
    "    data_dir=data_dir,\n",
    "    llm=llm,\n",
    "    chunk_size=DEFAULT_CHUNK_SIZE,\n",
    "    similarity_top_k=DEFAULT_TOP_K,\n",
    "    small_chunk_size=DEFAULT_SMALL_CHUNK_SIZE,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Pittsburgh can become a startup hub by leveraging its increasing population of young people, particularly those aged 25 to 29, who are crucial for the startup ecosystem. The city should encourage the youth-driven food boom, preserve historic buildings, and make the city more bicycle and pedestrian-friendly. Additionally, Carnegie Mellon University (CMU) should focus on being an even better research university to attract ambitious talent. The city should also foster a culture of tolerance and gradually build an investor community.\n",
       "\n",
       "There are two types of moderates: intentional moderates and accidental moderates. Intentional moderates deliberately choose positions midway between the extremes of right and left, while accidental moderates form their opinions independently on each issue, resulting in a broad range of views that average to a moderate position."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "# run a query\n",
    "res = await wf.run(\n",
    "    query_str=\"How can Pittsburgh become a startup hub, and what are the two types of moderates?\",\n",
    "    query_eng=result[\"query_engine\"],\n",
    ")\n",
    "display(Markdown(str(res)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-RvIdVF4_-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
